/* Copyright (C) 2010 Luca Piccinelli
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "generators/resource/D2DNBayesModelFileGenerator.h"
#include "machine_learning/statistical_models/normal_bayes_model.h"
#include "factory/ResourceGeneratorFactory.h"
#include "prj_definitions.h"

#include <string>

#include <cv.h>
#include <highgui.h>

using namespace std;
using namespace cv;

namespace NAMESPACE {

// ****** D2DNBayesModelFileGenerator implementation ****************************

//  _____________________________________________________
// |   Private functions                                 |
// |_____________________________________________________|

void D2DNBayesModelFileGenerator::free_element(IOElement* el){
    NormalBayesModel* model = boost::any_cast<NormalBayesModel*>( el->getElement() );
    Discret2DNormalBayesModel* d2d_model = dynamic_cast<Discret2DNormalBayesModel*>(model);
    delete d2d_model;
}

memory_size_t D2DNBayesModelFileGenerator::measure_memory(){
    if(!generated_buffer->size()) return 0;

    std::vector<IOElement*>::const_reverse_iterator it = generated_buffer->rbegin();
    NormalBayesModel* model = boost::any_cast<NormalBayesModel*>( (*it)->getElement() );
    Discret2DNormalBayesModel* d2d_model = dynamic_cast<Discret2DNormalBayesModel*>(model);

    return d2d_model->get_a_priori_data_lut().rows *
           d2d_model->get_a_priori_data_lut().cols *
           sizeof(prob_t)                          *
           (d2d_model->get_with_evidence_data_luts().size() + 1);
}

// .......................................................

//  _____________________________________________________
// |   Standard constructors                             |
// |_____________________________________________________|

D2DNBayesModelFileGenerator::~D2DNBayesModelFileGenerator(){
    free_buffer();
}
D2DNBayesModelFileGenerator::D2DNBayesModelFileGenerator() : AbstractFileGenerator(){
    key->setType(TypeParseTraits<NormalBayesModel*>::name());
}
D2DNBayesModelFileGenerator::D2DNBayesModelFileGenerator(const D2DNBayesModelFileGenerator& g)
                                                         : AbstractFileGenerator(g){
    key->setType(TypeParseTraits<NormalBayesModel*>::name());
}
D2DNBayesModelFileGenerator& D2DNBayesModelFileGenerator::operator=(const D2DNBayesModelFileGenerator& g){
    if(this == &g) return *this;
    AbstractFileGenerator::operator =(g);
    return *this;
}
// .......................................................

//  _____________________________________________________
// |   Others constructors                               |
// |_____________________________________________________|

D2DNBayesModelFileGenerator::D2DNBayesModelFileGenerator(const std::string& _resource,
                                                         const memory_size_t& _max_mem)
                                                         : AbstractFileGenerator(_resource, _max_mem){}
D2DNBayesModelFileGenerator::D2DNBayesModelFileGenerator(const IOKey& k,
                                                         const memory_size_t& _max_mem )
                                                         : AbstractFileGenerator(k, _max_mem){
    key->setType(TypeParseTraits<NormalBayesModel*>::name());
}
D2DNBayesModelFileGenerator::D2DNBayesModelFileGenerator(const std::string& _resource,
                                                         const IOKey& k,
                                                         const memory_size_t& _max_mem)
                                                         : AbstractFileGenerator(_resource, k, _max_mem){
    key->setType(TypeParseTraits<NormalBayesModel*>::name());
}
// .......................................................

Generator& D2DNBayesModelFileGenerator::set_key(const std::string& ns, const std::string& key_name){
    AbstractGenerator::set_key(ns, key_name);
    key->setType(TypeParseTraits<NormalBayesModel*>::name());
    return *this;
}

IOElement* D2DNBayesModelFileGenerator::generate(const IOKey& key){
    if(generated_buffer->size()) return (*generated_buffer->rbegin());

    FileStorage fs(resource, FileStorage::READ);
    FileNode a_posteriori_data_fn = fs[A_POSTERIORI_DATA_NODE_NAME];
    FileNode a_priori_data_fn     = fs[A_PRIORI_DATA_NODE_NAME];
    FileNode a_priori_class_fn    = fs[A_PRIORI_CLASS_NODE_NAME];

    FileNodeIterator fn_it = a_posteriori_data_fn.begin();
    vector<Mat>* a_posteriori_data = new vector<Mat>();
    int rows = 0, cols = 0;
    for(; fn_it != a_posteriori_data_fn.end(); fn_it++){
        rows = (*fn_it)["rows"];
        cols = (*fn_it)["cols"];
        Mat m;
        *fn_it >> m;
        a_posteriori_data->push_back(m);
    }
    Discret2DNormalBayesModel* model = new Discret2DNormalBayesModel(0, 0, rows, cols);
    model->set_with_evidence_data_luts(a_posteriori_data);

    a_priori_data_fn  >> model->get_a_priori_data_lut();
    a_priori_class_fn >> model->get_a_priori_evidence();


    NormalBayesModel* model_interface = model;
    IOElement* el = new HandledIOElement(model_interface, key);

    return el;
}

// -----------------------------------------------------------------------------

REGISTER_PARSE_TYPE(D2DNBayesModelFileGenerator);
ResourceGeneratorFactory<D2DNBayesModelFileGenerator> d2DNBayesModelFileGenerator_instance;
}